# -*- coding: utf-8 -*-
'''
Created on 23.07.2019

@author: yu03
'''

import cv2
import matplotlib.pyplot as plt
from scipy import signal
from FFT_Interpolation import FFT_interpolation_boxcar, FFT_interpolation_2, FFT_interpolation_compare, FFT_cal
import numpy as np
import os
from scipy.optimize import curve_fit


def running_mean(x, N):
    cumsum = np.cumsum(np.insert(x, 0, 0)) 
    return np.concatenate(( np.zeros(int((N)/2)), (cumsum[N:]-cumsum[:-N])/float(N) , np.zeros(int((N-1)/2)) ))
def fit_func(x, a, b, c):
    return a*(x-b)**2 + c

Lamda = 633e-9
pix_size = 5.3e-6
V_x, V_y, V_z = 0, 0, 0

img_set = []
hor_angle_centers = []
ver_angle_centers = []
hor_phase_centers = []
ver_phase_centers = []
hor_freq_set = []
ver_freq_set = []
hor_m_k_num_set = []
ver_m_k_num_set = []
hor_f_fit_set = []
ver_f_fit_set = []
hor_phi_fit_set = []
ver_phi_fit_set = []

for i in range(100):
    pattern_path = r'C:\Users\yu03\Desktop\Isolation Chamber Test 2019-07-19\No Glass Cover\Fixed Mirror_(ring far)_bmps\1\test_%i.bmp'%i
    img = cv2.imread(pattern_path, 0)
    hor_center = np.array(img[539])
    ver_center = img[:,662]
    hor_center = np.diff(hor_center.astype('float32'))
    hor_center = np.concatenate(([0], hor_center))
    ver_center = np.diff(ver_center.astype('float32'))
    ver_center = np.concatenate(([0], ver_center))
    
    DC_num = 0
    hor_freq_estim, hor_phase_estim, hor_freqline, hor_sig_magnitude, hor_sig_phase,  hor_m_k_num, hor_X_m_k, hor_freq_for_phase = FFT_interpolation_2(hor_center, pix_size, 1e5, DC_num)
    ver_freq_estim, ver_phase_estim, ver_freqline, ver_sig_magnitude, ver_sig_phase, ver_m_k_num, ver_X_m_k, ver_freq_for_phase = FFT_interpolation_2(ver_center, pix_size, 1e5, DC_num)
    
    hor_FFT_start = np.where(hor_sig_magnitude[DC_num:] > hor_X_m_k*0.4)[0][0]+DC_num
    hor_FFT_end = np.where(hor_sig_magnitude[DC_num:] > hor_X_m_k*0.4)[0][-1]+DC_num
    hor_fit_x = hor_freqline[hor_FFT_start:hor_FFT_end+1]
    hor_fit_y = hor_sig_magnitude[hor_FFT_start:hor_FFT_end+1]
    hor_fit_phase = hor_sig_phase[hor_FFT_start:hor_FFT_end+1]
    hor_params = curve_fit(fit_func, hor_fit_x, hor_fit_y)
    [hor_a, hor_b, hor_c] = hor_params[0]
    hor_f_fit = hor_b
    hor_f_fit_set.append(hor_f_fit)
#     hor_params = curve_fit(fit_func_phase, hor_fit_x, np.unwrap(hor_fit_phase))
#     [hor_q] = hor_params[0]
#     hor_phi_fit = p*hor_freqline[hor_m_k_num]+hor_q
#     hor_phi_fit_set.append(hor_phi_fit+4*np.pi-0.18)
    
    
    ver_FFT_start = np.where(ver_sig_magnitude[DC_num:] > ver_X_m_k*0.4)[0][0]+DC_num
    ver_FFT_end = np.where(ver_sig_magnitude[DC_num:] > ver_X_m_k*0.4)[0][-1]+DC_num
    ver_fit_x = ver_freqline[ver_FFT_start:ver_FFT_end+1]
    ver_fit_freq = ver_sig_magnitude[ver_FFT_start:ver_FFT_end+1]
    ver_params = curve_fit(fit_func, ver_fit_x, ver_fit_freq)
    [ver_a, ver_b, ver_c] = ver_params[0]
    ver_f_fit = ver_b
    ver_f_fit_set.append(ver_f_fit)
    
    print(i, hor_m_k_num, ver_m_k_num)
    hor_freq_set.append(hor_freq_estim)
    ver_freq_set.append(ver_freq_estim)
    
    hor_m_k_num_set.append(hor_m_k_num)
    ver_m_k_num_set.append(ver_m_k_num)
    
    hor_phase_centers.append(hor_phase_estim)
    ver_phase_centers.append(ver_phase_estim)
    
plt.figure(1)
plt.subplot(3,1,1)
plt.plot(hor_center)
plt.grid(which='major', axis='both')
plt.subplot(3,1,2)
plt.stem(hor_freqline, hor_sig_magnitude, use_line_collection=True)
plt.grid(which='major', axis='both')
plt.subplot(3,1,3)
plt.plot(ver_center)
plt.grid(which='major', axis='both')
plt.show()
 
plt.figure(2)
ax1 = plt.subplot(2,2,1)
plt.plot(hor_freq_set, color='blue', marker=' ')
plt.plot(hor_f_fit_set, color='red', marker=' ')
plt.title("Horizontal FFT")
plt.ylabel("Horizontal Freq Estimation (/m)")
plt.xlabel("Samples")
plt.grid(which='major', axis='both')
freq_range = plt.gca().get_ylim()
angle_range = ((V_x-Lamda*freq_range[0]/2)*1e6, (V_x-Lamda*freq_range[1]/2)*1e6)
ax1_angle = ax1.twinx()
plt.ylim(angle_range)
plt.ylabel('Horizontal Tilting (urad)')
  
ax2 = plt.subplot(2,2,2)
plt.plot(ver_freq_set, color='blue', marker=' ')
plt.plot(ver_f_fit_set, color='red', marker=' ')
plt.title("Vertical FFT")
plt.ylabel("Vertical Freq Estimation (/m)")
plt.xlabel("Samples")
plt.grid(which='major', axis='both')
freq_range = plt.gca().get_ylim()
angle_range = ((V_y-Lamda*freq_range[0]/2)*1e6, (V_y-Lamda*freq_range[1]/2)*1e6) 
ax2_angle = ax2.twinx()
plt.ylim(angle_range)
plt.ylabel('Vertical Tilting (urad)')

ax3 = plt.subplot(2,2,3)
hor_phase_unwrap = np.unwrap(hor_phase_centers)
# hor_phi_fit_set = np.unwrap(hor_phi_fit_set)
# hor_phase_unwrap = np.unwrap(hor_phase_centers)/2+np.pi/2
plt.plot(hor_phase_unwrap, color='blue', marker=' ')
# plt.plot(hor_phi_fit_set, color='red', marker=' ')
plt.title("Horizontal Phase")
plt.ylabel("Horizontal Phase (rad)")
plt.xlabel("Samples")
plt.grid(which='major', axis='both')
phase_range = plt.gca().get_ylim()
length_range = (phase_range[0]/4 /np.pi * Lamda*1e9, phase_range[1]/4 /np.pi * Lamda*1e9)
ax3_length = ax3.twinx()
plt.ylim(length_range)
plt.ylabel('Horizontal Length (nm)')

        
ax4 = plt.subplot(2,2,4)
ver_phase_unwrap = np.unwrap(ver_phase_centers)
# ver_phase_unwrap = np.unwrap(ver_phase_centers)/2+np.pi/2
plt.plot(ver_phase_unwrap, color='blue', marker=' ')
plt.title("Vertical Phase")
plt.ylabel("Vertical Phase (rad)")
plt.xlabel("Samples")
plt.grid(which='major', axis='both')
phase_range = plt.gca().get_ylim()
length_range = (phase_range[0]/4 /np.pi * Lamda*1e9, phase_range[1]/4 /np.pi * Lamda*1e9)
ax4_length = ax4.twinx()
plt.ylim(length_range)
plt.ylabel('Vertical Length (nm)')

plt.tight_layout()
     
      
plt.figure(3)
start_cutting, stop_cutting = 0, len(hor_phase_centers)
# start_cutting, stop_cutting = 100, 580
plt.subplot(2,2,1)
measured_phase = hor_phase_unwrap[start_cutting:stop_cutting]
linear_fit = np.linspace(measured_phase[0], measured_phase[-1], len(measured_phase))
nonlinearity = (measured_phase - linear_fit) / 4/ np.pi *Lamda
plt.plot(linear_fit, color='red', label='Linear Data')
plt.plot(measured_phase, color='blue', label='Measured Data')
legend = plt.legend(loc='upper right')
plt.setp(legend.get_texts()[0], color = 'red')
plt.setp(legend.get_texts()[1], color = 'blue')
plt.title("Horizontal Phase Measurement")
plt.ylabel("Horizontal Phase (rad)")
plt.xlabel("Frames")
plt.grid(which='major', axis='both')
plt.subplot(2,2,3)
plt.plot(nonlinearity*1e9)
plt.title("Horizontal Nonlinearity")
plt.ylabel("Displacement Nonlinearity (nm)")
plt.xlabel("Frames")
plt.grid(which='major', axis='both')
plt.subplot(2,2,2)
measured_phase = ver_phase_unwrap[start_cutting:stop_cutting]
linear_fit = np.linspace(measured_phase[0], measured_phase[-1], len(measured_phase))
nonlinearity = measured_phase - linear_fit
nonlinearity = (measured_phase - linear_fit) / 4/ np.pi *Lamda
plt.plot(linear_fit, color='red', label='Linear Data')
plt.plot(measured_phase, color='blue', label='Measured Data')
legend = plt.legend(loc='upper right')
plt.setp(legend.get_texts()[0], color = 'red')
plt.setp(legend.get_texts()[1], color = 'blue')
plt.title("Vertical Phase Measurement")
plt.ylabel("Vertical Phase (rad)")
plt.xlabel("Frames")
plt.grid(which='major', axis='both')
plt.subplot(2,2,4)
plt.plot(nonlinearity*1e9)
plt.title("Vertical Nonlinearity")
plt.ylabel("Displacement Nonlinearity (nm)")
plt.xlabel("Frames")
plt.grid(which='major', axis='both')
   
plt.tight_layout()
# length_freqline, length_FFT, length_magnitude, length_phase = FFT_cal(hor_phase_unwrap, 0.02)
# plt.figure(4)
# plt.stem(length_freqline, length_magnitude, linefmt='b', markerfmt='bo', use_line_collection='True')
plt.show()